#pragma once

#include <iostream>
#include "NvInfer.h"

namespace vortex
{
    class NvLogger : public nvinfer1::ILogger
    {
    public:
        const char* severity_string(nvinfer1::ILogger::Severity t)
        {
            switch(t){
                case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
                case nvinfer1::ILogger::Severity::kERROR:   return "error";
                case nvinfer1::ILogger::Severity::kWARNING: return "warning";
                case nvinfer1::ILogger::Severity::kINFO:    return "info";
                case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
                default: return "unknow";
            }
        }

        virtual void log(Severity severity, const char* msg) noexcept override
        {
            // suppress info-level message
            if (severity != Severity::kINFO)
            {
                std::cout << msg << std::endl;
            }
        }
    };
}
